#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : lei.wang
from typing import Union
from rest_framework.response import Response
from rest_framework import status


class BaseResponse(Response):

    def __init__(self, code=None, data=None, message='Success', **kwargs):
        result = {'code': code, 'data': data or {}, 'message': message}
        super().__init__(result, status=status.HTTP_200_OK, **kwargs)


def response_ok(data: Union[dict, list], message='success', **kwargs):
    return BaseResponse(data=data, code='20000', message=message, **kwargs)


def response_error(data: Union[dict, list], message='failed', **kwargs):
    return BaseResponse(data=data, code='20001', message=message, **kwargs)
